// License: Apache 2.0. See LICENSE file in root directory.
// Copyright(c) 2024 Intel Corporation. All Rights Reserved.
#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/CameraInfo.h>
#include <cv_bridge/cv_bridge.h>

#include <librealsense2/rs.hpp>
#include "librealsense2/rsutil.h"
#include "example-imgui.hpp"
#include <librealsense2-gl/rs_processing_gl.hpp> // Include GPU-Processing API
#include <opencv2/opencv.hpp>
#include <thread>

/*
 This example introduces the concept of spatial stream alignment.
 For example usecase of alignment, please check out align-advanced and measure demos.
 The need for spatial alignment (from here "align") arises from the fact
 that not all camera streams are captured from a single viewport.
 Align process lets the user translate images from one viewport to another.
 That said, results of align are synthetic streams, and suffer from several artifacts:
 1. Sampling - mapping stream to a different viewport will modify the resolution of the frame
               to match the resolution of target viewport. This will either cause downsampling or
               upsampling via interpolation. The interpolation used needs to be of type
               Nearest Neighbor to avoid introducing non-existing values.
 2. Occlussion - Some pixels in the resulting image correspond to 3D coordinates that the original
               sensor did not see, because these 3D points were occluded in the original viewport.
               Such pixels may hold invalid texture values.
*/

// This example assumes camera with depth and color
// streams, and direction lets you define the target stream
enum class direction
{
    to_depth,
    to_color
};

// Forward definition of UI rendering, implemented below
rs2_stream find_stream_to_align(const std::vector<rs2::stream_profile>& streams);

int main(int argc, char * argv[]) try
{   
    // 初始化ROS节点
    ros::init(argc, argv, "realsense_align_publisher");
    ros::NodeHandle nh;

    ros::Publisher color_pub = nh.advertise<sensor_msgs::Image>("/camera/color/image_raw", 10);
    ros::Publisher depth_pub = nh.advertise<sensor_msgs::Image>("/camera/aligned_depth_to_color/image_raw", 10);

    std::string serial;
    if (!device_with_streams({ RS2_STREAM_COLOR,RS2_STREAM_DEPTH }, serial))
        return EXIT_SUCCESS;

    bool use_gpu_processing = true;

    rs2::pipeline pipe;
    rs2::config cfg;
    if (!serial.empty())
        cfg.enable_device(serial);
    cfg.enable_stream(RS2_STREAM_COLOR,640, 480, RS2_FORMAT_RGB8, 60);
    cfg.enable_stream(RS2_STREAM_DEPTH,640, 480, RS2_FORMAT_Z16, 60);
    // pipe.start(cfg);
    cv::Mat imCV, depthCV;
    int width_img = 0, height_img = 0;
    double timestamp_image = -1.0;
    bool image_ready = false;
    int count_im_buffer = 0; // count dropped frames
    rs2::pipeline_profile pipe_profile = pipe.start(cfg);
    rs2_stream align_to = find_stream_to_align(pipe_profile.get_streams());
    rs2::gl::align align(align_to);

    std::cout << "align_to: " << align_to << std::endl;

    while (ros::ok()) // Application still alive?
    {
        // 等待获取帧集
        rs2::frameset fs;
        if (!pipe.poll_for_frames(&fs)) {
            // 如果没有可用的帧，继续等待
            std::this_thread::sleep_for(std::chrono::milliseconds(5));
            continue;
        }
        
        // 获取彩色和深度的帧流尺寸（第一次循环时设置）
        if (width_img == 0 || height_img == 0) {
            auto color = fs.get_color_frame();
            if (color) {
                width_img = color.get_width();
                height_img = color.get_height();
            } else {
                continue; // 等待有效的彩色帧
            }
        }
        
        // 处理对齐
        auto processed = align.process(fs);

        rs2::video_frame color_frame = processed.first(align_to);
        rs2::depth_frame depth_frame = processed.get_depth_frame();
        
        cv::Mat color_image = cv::Mat(cv::Size(width_img, height_img), CV_8UC3, (void*)(color_frame.get_data()), cv::Mat::AUTO_STEP);
            
        cv::Mat depth_image = cv::Mat(cv::Size(width_img, height_img), CV_16U, (void*)(depth_frame.get_data()), cv::Mat::AUTO_STEP);

        std_msgs::Header header;
        header.stamp = ros::Time::now();
        header.frame_id = "camera_link"; 
        
        sensor_msgs::ImagePtr color_msg = cv_bridge::CvImage(header, 
                                                           "rgb8", 
                                                           color_image).toImageMsg();
        color_pub.publish(color_msg);
        
        sensor_msgs::ImagePtr depth_msg = cv_bridge::CvImage(header, 
                                                           "16UC1", 
                                                           depth_image).toImageMsg();
        depth_pub.publish(depth_msg);
        
        // 处理ROS消息
        ros::spinOnce();
    }
    return EXIT_SUCCESS;
}
catch (const rs2::error & e)
{
    std::cerr << "RealSense error calling " << e.get_failed_function() << "(" << e.get_failed_args() << "):\n    " << e.what() << std::endl;
    return EXIT_FAILURE;
}
catch (const std::exception & e)
{
    std::cerr << e.what() << std::endl;
    return EXIT_FAILURE;
}

rs2_stream find_stream_to_align(const std::vector<rs2::stream_profile>& streams)
{
    //Given a vector of streams, we try to find a depth stream and another stream to align depth with.
    //We prioritize color streams to make the view look better.
    //If color is not available, we take another stream that (other than depth)
    rs2_stream align_to = RS2_STREAM_ANY;
    bool depth_stream_found = false;
    bool color_stream_found = false;
    for (rs2::stream_profile sp : streams)
    {
        rs2_stream profile_stream = sp.stream_type();
        if (profile_stream != RS2_STREAM_DEPTH)
        {
            if (!color_stream_found)         //Prefer color
                align_to = profile_stream;

            if (profile_stream == RS2_STREAM_COLOR)
            {
                color_stream_found = true;
            }
        }
        else
        {
            depth_stream_found = true;
        }
    }

    if(!depth_stream_found)
        throw std::runtime_error("No Depth stream available");

    if (align_to == RS2_STREAM_ANY)
        throw std::runtime_error("No stream found to align with Depth");

    return align_to;
}